#ifdef __aarch64__

.text
.align 5
.global IndirectGemmInt8_4x4_pc
#ifndef __APPLE__
.type IndirectGemmInt8_4x4_pc, %function
#endif

// void IndirectGemmInt8_4x4_pc(int8_t *output, int16_t *input, int16_t *weight, int16_t *bias, size_t k4, size_t ic,
// size_t oc, size_t offset, int32_t *act_min, int32_t *act_max, int32_t *out_zp, int32_t *out_multiplier,
// int32_t *shift_before, int32_t *shift_after, size_t filter_per_channel, size_t output_per_channel);
// x0: output, x1: input, x2: weight, x3: bias, x4: k4, x5: ic, x6: oc, x7: offset
// #0: act_min, #8: act_max, #16: out_zp, #24: multiplier, #32: shift_before, #40: shift_after
// we use sdot intrinsic on cores that supports dotprod(Armv8.2-A w/dp or later)
// mrs intrinsic could read system register ID_AA64ISAR0_EL1(or s3_0_c0_c6_0 on Armv8.2-A)
// the 44-48 bits indicates whether dotprod is supported
IndirectGemmInt8_4x4_pc:

    .macro INIT_RES
        dup v28.4s, wzr
        dup v29.4s, wzr
        dup v30.4s, wzr
        dup v31.4s, wzr
    .endm

    .macro INIT_SUM
        dup v8.4s, wzr
        dup v9.4s, wzr
        dup v10.4s, wzr
        dup v11.4s, wzr
        dup v12.4s, wzr
        dup v13.4s, wzr
        dup v14.4s, wzr
        dup v15.4s, wzr
        dup v16.4s, wzr
        dup v17.4s, wzr
        dup v18.4s, wzr
        dup v19.4s, wzr
        dup v20.4s, wzr
        dup v21.4s, wzr
        dup v22.4s, wzr
        dup v23.4s, wzr
    .endm

    // registers v8 ~ v15 must be preserved by a callee across subroutine calls, according to
    // https://github.com/ARM-software/abi-aa/blob/master/aapcs64/aapcs64.rst#simd-and-floating-point-registers
    // r19 ~ r29 should be also preserved
    // whereas our coding style do not permit such amount of parameters
    sub sp, sp, #176
    st1 {v8.4s, v9.4s, v10.4s, v11.4s}, [sp], #64
    st1 {v12.4s, v13.4s, v14.4s, v15.4s}, [sp], #64
    stp x19, x20, [sp], #16
    stp x21, x22, [sp], #16
    stp x23, x24, [sp], #16
    stp x25, x26, [sp], #16

    ldr x8, [sp, #0]
    ldr x9, [sp, #8]
    ldr x16, [sp, #16]
    ldr x17, [sp, #24]
    ldr x18, [sp, #32]
    ldr x19, [sp, #40]
    ldr x20, [sp, #48]
    ldr x21, [sp, #56]

    mov x15, x17
    mov x25, x18
    mov x26, x19
    add x16, x16, #4

    LoopOc:
        mov x10, x5
        mov x12, x1
        mov x11, x0
        ld1r {v7.4s}, [x3], #4

        cbz x20, NoQuantizeReset
        mov x22, x17
        mov x23, x18
        mov x24, x19
    NoQuantizeReset:
        INIT_RES

        LoopIc:
            INIT_SUM
            mov x13, x4

            LoopKsize:
                // load input for output 1-8
                ld1r {v0.2d}, [x12], #8
                ld1r {v1.2d}, [x12], #8
                ld1r {v2.2d}, [x12], #8
                ld1r {v3.2d}, [x12], #8
                // load weight
                ld1 {v4.8h, v5.8h}, [x2], #32
                // step for output 1-4
                smlal v8.4s, v4.4h, v0.4h
                smlal2 v9.4s, v4.8h, v0.8h
                smlal v10.4s, v5.4h, v0.4h
                smlal2 v11.4s, v5.8h, v0.8h
                smlal v12.4s, v4.4h, v1.4h
                smlal2 v13.4s, v4.8h, v1.8h
                smlal v14.4s, v5.4h, v1.4h
                smlal2 v15.4s, v5.8h, v1.8h
                smlal v16.4s, v4.4h, v2.4h
                smlal2 v17.4s, v4.8h, v2.8h
                smlal v18.4s, v5.4h, v2.4h
                smlal2 v19.4s, v5.8h, v2.8h
                smlal v20.4s, v4.4h, v3.4h
                smlal2 v21.4s, v4.8h, v3.8h
                smlal v22.4s, v5.4h, v3.4h
                smlal2 v23.4s, v5.8h, v3.8h

                subs x13, x13, #1
                bne LoopKsize

            addp v8.4s, v8.4s, v9.4s
            addp v10.4s, v10.4s, v11.4s
            addp v12.4s, v12.4s, v13.4s
            addp v14.4s, v14.4s, v15.4s
            addp v16.4s, v16.4s, v17.4s
            addp v18.4s, v18.4s, v19.4s
            addp v20.4s, v20.4s, v21.4s
            addp v22.4s, v22.4s, v23.4s

            addp v8.4s, v8.4s, v10.4s
            addp v12.4s, v12.4s, v14.4s
            addp v16.4s, v16.4s, v18.4s
            addp v20.4s, v20.4s, v22.4s

            ld1r {v1.4s}, [x23], #4
            ld1r {v2.4s}, [x22], #4
            ld1r {v3.4s}, [x24], #4

            sqshl v8.4s, v8.4s, v1.4s
            sqshl v12.4s, v12.4s, v1.4s
            sqshl v16.4s, v16.4s, v1.4s
            sqshl v20.4s, v20.4s, v1.4s

            sqrdmulh v8.4s, v8.4s, v2.4s
            sqrdmulh v12.4s, v12.4s, v2.4s
            sqrdmulh v16.4s, v16.4s, v2.4s
            sqrdmulh v20.4s, v20.4s, v2.4s

            and v0.16b, v3.16b, v8.16b
            sshr v0.4s, v0.4s, #31
            sqadd v8.4s, v8.4s, v0.4s
            srshl v8.4s, v8.4s, v3.4s
            and v0.16b, v3.16b, v12.16b
            sshr v0.4s, v0.4s, #31
            sqadd v12.4s, v12.4s, v0.4s
            srshl v12.4s, v12.4s, v3.4s
            and v0.16b, v3.16b, v16.16b
            sshr v0.4s, v0.4s, #31
            sqadd v16.4s, v16.4s, v0.4s
            srshl v16.4s, v16.4s, v3.4s
            and v0.16b, v3.16b, v20.16b
            sshr v0.4s, v0.4s, #31
            sqadd v20.4s, v20.4s, v0.4s
            srshl v20.4s, v20.4s, v3.4s

            add v28.4s, v28.4s, v8.4s
            add v29.4s, v29.4s, v12.4s
            add v30.4s, v30.4s, v16.4s
            add v31.4s, v31.4s, v20.4s

            subs x10, x10, #1
            bne LoopIc

        Activation:
            cbz x20, NoPerOC
            cbz x21, NoPerOC
            ld1r {v0.4s}, [x25], #4
            ld1r {v1.4s}, [x15], #4
            ld1r {v2.4s}, [x26], #4
            ld1r {v3.4s}, [x8], #4
            ld1r {v4.4s}, [x9], #4
            ld1r {v5.4s}, [x16]
            add x16, x16, #8
            b QuantizeStart
        NoPerOC:
            ld1r {v0.4s}, [x25]
            ld1r {v1.4s}, [x15]
            ld1r {v2.4s}, [x26]
            ld1r {v3.4s}, [x8]
            ld1r {v4.4s}, [x9]
            ld1r {v5.4s}, [x16]
        QuantizeStart:
            sqshl v6.4s, v7.4s, v0.4s
            sqrdmulh v6.4s, v6.4s, v1.4s
            and v0.16b, v2.16b, v6.16b
            sshr v0.4s, v0.4s, #31
            sqadd v6.4s, v6.4s, v0.4s
            srshl v6.4s, v6.4s, v2.4s

            add v28.4s, v28.4s, v6.4s
            add v29.4s, v29.4s, v6.4s
            add v30.4s, v30.4s, v6.4s
            add v31.4s, v31.4s, v6.4s

            smin v28.4s, v28.4s, v4.4s
            smin v29.4s, v29.4s, v4.4s
            smin v30.4s, v30.4s, v4.4s
            smin v31.4s, v31.4s, v4.4s

            smax v28.4s, v28.4s, v3.4s
            smax v29.4s, v29.4s, v3.4s
            smax v30.4s, v30.4s, v3.4s
            smax v31.4s, v31.4s, v3.4s

            sqxtn v0.4h, v28.4s
            sqxtn2 v0.8h, v29.4s
            sqxtn v6.8b, v2.8h
            sqxtn v1.4h, v30.4s
            sqxtn2 v1.8h, v31.4s
            sqxtn2 v6.16b, v1.8h

        // prefetching is not prefered while writing results in spite of cache missings
        // you could try prfm pstl2strm
        WriteStart:
            cmp x6, #1
            beq Write1
            cmp x6, #2
            beq Write2
            cmp x6, #3
            beq Write3
            b Write4
        Write1:
            st1 {v6.b}[0], [x11], x7
            st1 {v6.b}[4], [x11], x7
            st1 {v6.b}[8], [x11], x7
            st1 {v6.b}[12], [x11], x7
            add x0, x0, #1
            b WriteEnd
        Write2:
            st1 {v6.h}[0], [x11], x7
            st1 {v6.h}[2], [x11], x7
            st1 {v6.h}[4], [x11], x7
            st1 {v6.h}[6], [x11], x7
            add x0, x0, #2
            b WriteEnd
        Write3:
            add x14, x11, #2
            st1 {v6.h}[0], [x11], x7
            st1 {v6.b}[2], [x14], x7
            st1 {v6.h}[2], [x11], x7
            st1 {v6.b}[6], [x14], x7
            st1 {v6.h}[4], [x11], x7
            st1 {v6.b}[10], [x14], x7
            st1 {v6.h}[6], [x11], x7
            st1 {v6.b}[14], [x14], x7
            add x0, x0, #3
            b WriteEnd
        Write4:
            st1 {v6.s}[0], [x11], x7
            st1 {v6.s}[1], [x11], x7
            st1 {v6.s}[2], [x11], x7
            st1 {v6.s}[3], [x11], x7
            add x0, x0, #4

        subs x6, x6, #4
        add x3, x3, #16
        bgt LoopOc

    sub sp, sp, #176
    ld1 {v8.4s, v9.4s, v10.4s, v11.4s}, [sp], #64
    ld1 {v12.4s, v13.4s, v14.4s, v15.4s}, [sp], #64
    ldp x19, x20, [sp], #16
    ldp x21, x22, [sp], #16
    ldp x23, x24, [sp], #16
    ldp x25, x26, [sp], #16
    ret
#endif
